#!/bin/bash

# ----------------------------------------------------------------------
# Filename:   Partition_parted.sh
# Version:    1.0
# Date:       2013/12/06
# Author:     yuanhui.shi
# Email:      yuanhui.shi@cs2c.com.cn
# Summary:    Create a partition and delete a partition by parted command
# Notes:      Create a extended partition and delete a partition
# Copyright:  China Standard Software Co., Ltd.
# History：     
#             Version 1.0, 2013/12/06
#             - The first one
#             Version 1.1, 2014/04/01
#             - add CreatePrimary function
# ----------------------------------------------------------------------
. ${SFROOT}/lib/Echo.sh


##! @TODO: Create a extended partition when nothing extended partition exists.
##! @AUTHOR: yuanhui.shi
##! @VERSION: 1.0 
##! @OUT: Add a extended partition.

function CreateExtended ()
{
    parted -s /dev/sda mkpart extended ${STARTBLOCK}GB ${TOTAL}GB
    EchoResult "create extended"

    # get created extend partition, through grep extended and tail line, check extended partition
    EXTENDED_PART=`parted -s /dev/sda print | grep extended | awk '{print $1}' |tr -d ''`
    [ "x" != x${EXTENDED_PART} -a ${EXTENDED_PART} = `parted -s /dev/sda print |tail -2 | awk '/./{print $1}'` ]
    EchoResult "get created extended partition: ${EXTENDED_PART}"

    CreateLogical
}

##! @TODO: Create a logical partition when extended partition has existed.
##! @AUTHOR: yuanhui.shi
##! @VERSION: 1.0 
##! @OUT: Add a logical partition.

function CreateLogical ()
{
# from last partition END ~ END+5 or END ~ TOTAL

    ENDBLOCK=`echo "${STARTBLOCK}" | awk '{print $1+5}'`
    if [ $(echo "$ENDBLOCK > $TOTAL" | bc ) = 1 ];then
        ENDBLOCK=${TOTAL}
    fi

    parted -s /dev/sda mkpart logic ${STARTBLOCK}GB ${ENDBLOCK}GB  
    EchoResult "create logiacl"

    # get created logical partition
    LOGICAL_PART=`parted -s /dev/sda print |tail -2 | awk '/./{print $1}'`
    if [ "x" != x${LOGICAL_PART} ];then
        EchoInfo "get created logiacl partition: ${LOGICAL_PART}"
        return ${LOGICAL_PART}
    else
        echo  -e "\033[31m FAIL\033[0m | get created logical partition"
        return 0
    fi

}

##! @TODO: Create a primary partition in gpt.
##! @AUTHOR: huijing.hei
##! @VERSION: 1.0 
##! @OUT: Add a primary partition.

function CreatePrimary ()
{
# from last partition END ~ END+5 or END ~ TOTAL

    ENDBLOCK=`echo "${STARTBLOCK}" | awk '{print $1+5}'`
    if [ $(echo "$ENDBLOCK > $TOTAL" | bc ) = 1 ];then
        ENDBLOCK=${TOTAL}
    fi

    parted -s /dev/sda mkpart primary ${STARTBLOCK}GB ${ENDBLOCK}GB  
    EchoResult "create primary"

    # get created primary partition
    PRIMARY_PART=`parted -s /dev/sda print |tail -2 | awk '/./{print $1}'`
    if [ "x" != x${PRIMARY_PART} ];then
        EchoInfo "get created primary partition: ${PRIMARY_PART}"
        return ${PRIMARY_PART}
    else
        echo  -e "\033[31m FAIL\033[0m | get primary logical partition"
        return 0
    fi

}

##! @TODO: Create a extended partition in different conditions
##! @AUTHOR: yuanhui.shi
##! @VERSION: 1.0 
##! @OUT: Add a extended partition.

function CreatePartition ()
{
    HaveExtended=2

# check total block, is not null
    declare TOTAL=`parted -s /dev/sda print |grep Disk | awk -F ":" '{print $2}' | tr -d " GB"`
    if [ "x" = "x${TOTAL}" ] ;then
        echo -e "\033[31m FAIL\033[0m | get total block of /dev/sda"
        exit 1
    fi

# check STARTBLOCK, if greate than TOTAL BLOCK, then fail
    declare STARTBLOCK=`parted -s /dev/sda print |tail -2 | awk '/./{print $3}'| tr -d "GB"`
    if [ "x" = "x${STARTBLOCK}" ] ;then
        echo -e "\033[31m FAIL\033[0m | get start block of /dev/sda"
        exit 1
    elif [ $(echo "$STARTBLOCK >= $TOTAL" | bc ) = 1 ];then
        echo -e "\033[31m FAIL\033[0m | disk have not enough space to mkpart"
        exit 1
    fi

# get part table types, msdos or gpt
    PART_TABLE=`parted /dev/sda print | grep -i "Partition Table" | awk -F ":" '{print $2}' | tr -d ' ' | tr "A-Z" "a-z"`

    PRIMARY=`parted /dev/sda print | grep primary | wc -l`

    # check msdos
    if [ ${PART_TABLE} = "msdos" ];then

        # already has 4 primary partitions
        if [ $PRIMARY -ge 4 ];then
            HaveExtended=2
            EchoInfo "${PART_TABLE}: already have 4 primary partitions"
            exit 1
        # check extended partition
        elif [ -z "`parted -s /dev/sda print |grep extended`" ];then

            HaveExtended=1
            CreateExtended
        else
            HaveExtended=0
            CreateLogical
            return $?
        fi

    # check gpt
    elif  [ ${PART_TABLE} = "gpt" ];then

        # already has 128 primary partitions
        if [ $PRIMARY -ge 128 ];then
            HaveExtended=2
            EchoInfo "${PART_TABLE}: already have 128 primary partitions"
            exit 1
        else
            HaveExtended=3
            CreatePrimary
            return $?
        fi

    fi

  
#  parted -s /dev/sda print &> /tmp/tmp2
## check /tmp/tmp2
#    if [ ! -s /tmp/tmp2 ];then
#        echo -e "\033[31m FAIL\033[0m | get file /tmp/tmp2 is zero"
#    fi
#
#  if [ -n "`diff -b /tmp/tmp1 /tmp/tmp2 |grep Extended`" -a "`diff -b /tmp/tmp1 /tmp/tmp2 |sed '1d'|wc -l`" -le 1 ];then
#     echo -e "\033[31m FAIL\033[0m | 新建分区失败"
#     exit 1
#  elif [ -z "`diff -b /tmp/tmp1 /tmp/tmp2`" ];then
#     echo -e "\033[31m FAIL\033[0m | 新建分区失败"
#     exit 1
#  else echo -e  "        \033[32m PASS\033[0m | 新建分区成功"
#  fi
}

##! @TODO: Delete a logical partition.
##! @AUTHOR: yuanhui.shi
##! @VERSION: 1.0 
##! @OUT: Delete a logical partition.

function DeleteLogical ()
{
#    NUM1=`parted -s /dev/sda print | tail -2 |cut -d " " -f2`
    if [ "x" != x${LOGICAL_PART} ];then
        EchoInfo "to be deleted logical partition: ${LOGICAL_PART}"
    else
        echo  -e "\033[31m FAIL\033[0m | get logical partition to be deleted"
        exit 1
    fi

    parted -s /dev/sda rm $LOGICAL_PART
    EchoResult "rm partition $LOGICAL_PART"
#    return 0
}

##! @TODO: Delete a extended.
##! @AUTHOR: yuanhui.shi
##! @VERSION: 1.0 
##! @OUT: Delete the extended partition.

function DeleteExtended ()
{
#    NUM1=`diff -b /tmp/tmp1 /tmp/tmp2 | sed '1d'|cut -d " " -f3 | tac`
    if [ "x" != x${EXTENDED_PART} ];then
        EchoInfo "to be deleted logical partition: ${EXTENDED_PART}"
    else
        echo  -e "\033[31m FAIL\033[0m | get logical partition to be deleted"
        exit 1
    fi

    for i in $LOGICAL_PART $EXTENDED_PART
     do 
       {
         parted -s /dev/sda rm $i &>/dev/null
         EchoResult "rm partition $i"
       }
   done
#    return 0
}

##! @TODO: Delete a primary.
##! @AUTHOR: huijing.hei
##! @VERSION: 1.0 
##! @OUT: Delete the primary partition.

function DeletePrimary ()
{
#    NUM1=`diff -b /tmp/tmp1 /tmp/tmp2 | sed '1d'|cut -d " " -f3 | tac`
    if [ "x" != x${PRIMARY_PART} ];then
        EchoInfo "to be deleted primary partition: ${PRIMARY_PART}"
    else
        echo  -e "\033[31m FAIL\033[0m | get primary partition to be deleted"
        exit 1
    fi

    for i in $PRIMARY_PART
     do 
       {
         parted -s /dev/sda rm $i &>/dev/null
         EchoResult "rm partition $i"
       }
   done
#    return 0
}

##! @TODO: Delete a partition.
##! @AUTHOR: yuanhui.shi
##! @VERSION: 1.0 
##! @OUT: Delete the extended partition.

function DeletePartition ()
{  

#   if [ -n "`diff -b /tmp/tmp1 /tmp/tmp2`" ];then
      
#       {
        if [ "$HaveExtended" = 1 ] ;then 
               DeleteExtended
        elif [ "$HaveExtended" = 0 ] ;then
               DeleteLogical
        elif [ "$HaveExtended" = 3 ] ;then
               DeletePrimary
        elif  [ "$HaveExtended" = 2 ] ;then
               EchoInfo "no partition to be deleted"
        fi
#        }
#   fi     
   return 0
}
